import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# 准备的测试数据集
test_data = torchvision.datasets.CIFAR10("./dataset", train = False, transform = torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset = test_data, batch_size = 64, shuffle = True, num_workers = 0, drop_last = True)
# 上述的dataloader中参数batch_size = 64, 则dataloader会将64个图片和4个target分别打包，得到一个四维的tensor以及长度为4的一维tensor
# 如果dataloader中drop_last = True，则会将最后一个不足64个的batch丢弃
# 如果dataloader中shuffle = True，则会将数据集打乱
# 如果dataloader中num_workers = 0，则会使用主进程加载数据集，如果num_workers > 0，则会使用多个进程加载数据集


# 测试数据集中第一张图片及target
img, target = test_data[0]
print(img.shape)
print(target)

writer = SummaryWriter("dataloader")

for epoch in range(2):
    step = 0
    for data in test_loader: # data是一个batch的数据
        imgs, targets = data # imgs是一个四维的tensor，在这个例子中是64张图片，每张图片是3通道，高度，宽度。targets是一个一维的tensor，长度为64
        writer.add_images("Epoch:{}".format(epoch), imgs, step)
        step += 1

writer.close()